import json

from algorithm.antAlgorithm.ant.aco import Aco
from common.common_fun import initData


def aco_result(data):
    (dim_n, dim_k, dim_c, start_node, goal_node, is_key, edges, cost_v, cost_e,
     node_index_dict, node_edge_dict) = initData(data)
    cost_e = [[inner_element[0] for inner_element in sublist] for sublist in cost_e]
    print(cost_e)

    num_ants = 10
    num_iterations = 100
    alpha = 1
    beta = 5
    rho = 0.5
    Q = 1000

    solver = Aco(is_key, node_edge_dict, cost_e, start_node, goal_node, num_ants, num_iterations, alpha, beta, rho, Q)
    result, display = solver.work()

    out_map = dict()
    in_map = display["pathMap"]

    for (outer_key, inner_key), (value, lst) in in_map.items():
        # 如果outer_key尚未出现在新字典中，则添加它并初始化为空字典
        if outer_key not in out_map:
            out_map[node_index_dict.get_key(outer_key)] = {}

        # 将(outer_key, inner_key)对转换为嵌套字典的结构
        out_map[node_index_dict.get_key(outer_key)][node_index_dict.get_key(inner_key)] = value

    display["pathMap"] = out_map

    for ants in display["iteration_results"]:
        for ant in ants["ants"]:
            for i in range(0, len(ant["path"]["node"])):
                ant["path"]["node"][i] = node_index_dict.get_key(ant["path"]["node"][i])

    print(result)
    print(json.dumps(display))
    path_node = [node_index_dict.get_key(node) for node in result]
    path_edge = list()
    for i in range(len(result) - 1):
        path_edge.append(node_edge_dict[(result[i], result[i + 1])])
    return path_node, path_edge


def aco_display(data):
    # 初始化数据
    (dim_n, dim_k, dim_c, start_node, goal_node, is_key, edges, cost_v, cost_e,
     node_index_dict, node_edge_dict) = initData(data)
    cost_e = [[inner_element for inner_element in sublist] for sublist in cost_e]

    num_ants = data["antNum"]
    num_iterations = data["iterationNum"]
    alpha = data["a"]
    beta = data["b"]
    rho = data["r"]
    Q = data["Q"]
    # 计算
    solver = Aco(is_key, node_edge_dict, cost_e, start_node, goal_node, num_ants, num_iterations, alpha, beta, rho, Q)
    result, display = solver.work()
    out_map = dict()
    in_map = display["pathMap"]

    for (outer_key, inner_key), (value, lst) in in_map.items():
        # 如果outer_key尚未出现在新字典中，则添加它并初始化为空字典
        if node_index_dict.get(outer_key) not in out_map:
            out_map[node_index_dict.get(outer_key)] = {}

        # 将(outer_key, inner_key)对转换为嵌套字典的结构
        out_map[node_index_dict.get(outer_key)][node_index_dict.get(inner_key)] = value

    display["pathMap"] = out_map
    for ants in display["iteration_results"]:
        for ant in ants["ants"]:
            for i in range(0, len(ant["path"]["node"])):
                ant["path"]["node"][i] = node_index_dict.get(ant["path"]["node"][i])



    return display
